from PyQt5.QtWidgets import (QApplication, QLabel, QWidget, QVBoxLayout,
                             QPushButton, QTextEdit, QHBoxLayout, QFileDialog)
from PyQt5.QtGui import (QPixmap, QPainter, QColor, QFont, QPen,
                         QImage, QScreen, QCursor)
from PyQt5.QtCore import Qt, QRect, QPoint, pyqtSignal


class EditableScreenshot(QWidget):
    textAdded = pyqtSignal(str, QPoint)

    def __init__(self, pixmap=None, parent=None):
        super().__init__(parent)
        self.pixmap = pixmap
        self.drawing_text = False
        self.text_position = QPoint()
        self.text = ""
        self.font = QFont("Arial", 14)
        self.text_color = QColor(255, 255, 255)

    def setPixmap(self, pixmap):
        self.pixmap = pixmap
        self.update()

    def startTextDrawing(self):
        self.drawing_text = True
        self.setCursor(Qt.CrossCursor)

    def setTextProperties(self, text, font=None, color=None):
        self.text = text
        if font:
            self.font = font
        if color:
            self.text_color = color
        self.update()

    def mousePressEvent(self, event):
        if self.drawing_text and event.button() == Qt.LeftButton:
            self.text_position = event.pos()
            self.textAdded.emit(self.text, self.text_position)
            self.drawing_text = False
            self.setCursor(Qt.ArrowCursor)

    def paintEvent(self, event):
        if not self.pixmap:
            return

        painter = QPainter(self)
        painter.drawPixmap(0, 0, self.pixmap.scaled(
            self.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))

        if self.text and self.text_position:
            painter.setFont(self.font)
            painter.setPen(self.text_color)

            # 计算缩放比例
            scale_x = self.pixmap.width() / self.width()
            scale_y = self.pixmap.height() / self.height()

            # 转换坐标到原始图像坐标
            original_pos = QPoint(
                int(self.text_position.x() * scale_x),
                int(self.text_position.y() * scale_y)
            )

            text_image = QImage(self.pixmap.size(), QImage.Format_ARGB32)
            text_image.fill(Qt.transparent)

            text_painter = QPainter(text_image)
            text_painter.setFont(self.font)
            text_painter.setPen(self.text_color)

            text_rect = QRect(original_pos, QPoint())
            text_rect = text_painter.boundingRect(
                text_rect, Qt.AlignLeft | Qt.AlignTop, self.text)

            # 绘制文字背景（可选）
            text_painter.fillRect(
                text_rect.adjusted(-2, -2, 2, 2),
                QColor(0, 0, 0, 150))

            text_painter.drawText(text_rect, self.text)
            text_painter.end()

            painter.drawImage(0, 0, text_image.scaled(
                self.size(), Qt.KeepAspectRatio, Qt.SmoothTransformation))


class ScreenshotApp(QWidget):
    def __init__(self):
        super().__init__()
        self.initUI()

    def initUI(self):
        self.setWindowTitle("截图文字编辑")
        self.setGeometry(100, 100, 800, 600)

        self.editor = EditableScreenshot()
        self.text_edit = QTextEdit()
        self.text_edit.setPlaceholderText("输入要添加的文字...")

        self.btn_capture = QPushButton("截图")
        self.btn_add_text = QPushButton("添加文字")
        self.btn_save = QPushButton("保存")

        layout = QVBoxLayout()
        layout.addWidget(self.editor)

        text_layout = QHBoxLayout()
        text_layout.addWidget(self.text_edit)
        text_layout.addWidget(self.btn_add_text)
        layout.addLayout(text_layout)

        btn_layout = QHBoxLayout()
        btn_layout.addWidget(self.btn_capture)
        btn_layout.addWidget(self.btn_save)
        layout.addLayout(btn_layout)

        self.setLayout(layout)

        self.btn_capture.clicked.connect(self.capture_screen)
        self.btn_add_text.clicked.connect(self.prepare_add_text)
        self.btn_save.clicked.connect(self.save_screenshot)
        self.editor.textAdded.connect(self.add_text_to_image)

        self.btn_add_text.setEnabled(False)
        self.btn_save.setEnabled(False)

    def capture_screen(self):
        self.hide()
        QApplication.processEvents()

        screen = QApplication.primaryScreen()
        pixmap = screen.grabWindow(0)

        self.show()
        self.editor.setPixmap(pixmap)
        self.btn_add_text.setEnabled(True)
        self.btn_save.setEnabled(True)

    def prepare_add_text(self):
        self.editor.startTextDrawing()

    def add_text_to_image(self, text, position):
        self.editor.setTextProperties(self.text_edit.toPlainText())

    def save_screenshot(self):
        if not self.editor.pixmap:
            return

        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存截图", "", "PNG Files (*.png);;JPEG Files (*.jpg *.jpeg)")

        if file_path:
            # 创建一个新的QPixmap用于保存
            pixmap = self.editor.pixmap.copy()
            painter = QPainter(pixmap)

            # 添加所有文本到原始图像
            if self.editor.text:
                painter.setFont(self.editor.font)
                painter.setPen(self.editor.text_color)

                # 计算文本位置（需要转换回原始坐标）
                scale_x = pixmap.width() / self.editor.width()
                scale_y = pixmap.height() / self.editor.height()
                pos = QPoint(
                    int(self.editor.text_position.x() * scale_x),
                    int(self.editor.text_position.y() * scale_y)
                )

                text_rect = QRect(pos, QPoint())
                text_rect = painter.boundingRect(
                    text_rect, Qt.AlignLeft | Qt.AlignTop, self.editor.text)

                painter.fillRect(
                    text_rect.adjusted(-2, -2, 2, 2),
                    QColor(0, 0, 0, 150))

                painter.drawText(text_rect, self.editor.text)

            painter.end()

            if file_path.lower().endswith('.png'):
                pixmap.save(file_path, "PNG")
            else:
                pixmap.save(file_path, "JPG", quality=90)


if __name__ == "__main__":
    app = QApplication([])
    window = ScreenshotApp()
    window.show()
    app.exec_()